package aes

import (
	"crypto/aes"
	"crypto/cipher"
	"io"

	"servers/crypto/internal/utils"
)

// NewCipher key要求16, 24, 32位, iv 的长度必须等于block的长度
func NewCipher(key []byte) *CipherAES {
	block, err := aes.NewCipher(key)
	if err != nil {
		panic(err)
	}
	r := &CipherAES{
		block: block,
	}
	return r
}

type CipherAES struct {
	block cipher.Block
}

func (a *CipherAES) BlockSize() int {
	return a.block.BlockSize()
}

func (a *CipherAES) Encrypt(src, dst []byte) ([]byte, error) {
	bs := a.block.BlockSize()
	origData := utils.PaddingBlock(src, bs)
	if len(dst) < len(origData) {
		return nil, io.ErrShortWrite
	}
	out := dst[:len(origData)]
	for len(origData) > 0 {
		a.block.Encrypt(dst, origData[:bs])
		origData = origData[bs:]
		dst = dst[bs:]
	}
	return out, nil
}

func (a *CipherAES) Decrypt(src, dst []byte) ([]byte, error) {
	bs := a.block.BlockSize()
	if len(src)%bs != 0 {
		return nil, io.ErrUnexpectedEOF
	}
	if len(dst) < len(src) {
		return nil, io.ErrShortWrite
	}
	out := dst[:len(src)]
	for len(src) > 0 {
		a.block.Decrypt(dst, src[:bs])
		src = src[bs:]
		dst = dst[bs:]
	}
	out = utils.TrimPaddedBlock(out)
	return out, nil
}
